#!/usr/bin/env python

import numpy
import scipy.linalg
import copy
import sys

zero = numpy.float64(0.0)
one = numpy.float64(1.0)
two = numpy.float64(2.0)
half = numpy.float64(0.5)


#-----------------------------------------------------------------------------

def check_matrix_type(Mlist):
   for M in Mlist:
      assert isinstance(M, numpy.matrix)


#-----------------------------------------------------------------------------

class Two_site_fci_basis:

   def __init__(self):
      self.dets = []
      det = Determinant(nsites,nalpha,nbeta)
      det.weight = zero
      det.phiA[0,0] = 1
      det.phiB[0,0] = 1
      self.dets.append(det)
      #print "basis state 0:\n", det
      det = Determinant(nsites,nalpha,nbeta)
      det.weight = zero
      det.phiA[0,0] = 1
      det.phiB[1,0] = 1
      self.dets.append(det)
      #print "basis state 1:\n", det
      det = Determinant(nsites,nalpha,nbeta)
      det.weight = zero
      det.phiA[1,0] = 1
      det.phiB[0,0] = 1
      self.dets.append(det)
      #print "basis state 2:\n", det
      det = Determinant(nsites,nalpha,nbeta)
      det.weight = zero
      det.phiA[1,0] = 1
      det.phiB[1,0] = 1
      self.dets.append(det)
      #print "basis state 3:\n", det
      

   def project_into_basis(self,dets):

      # Accumulate basis fnc coefficients
      for basis_fnc in self.dets:
         basis_fnc.weight = zero
         for det in dets:
             wt = det.weight * det.get_overlap_integrals(basis_fnc)[0]
             basis_fnc.weight = basis_fnc.weight + wt

      # Normalize
##      wt = zero
##      for basis_fnc in self.dets:
##         wt = wt + basis_fnc.weight*basis_fnc.weight
##      wt = numpy.sqrt(wt)
##      for basis_fnc in self.dets:
##         basis_fnc.weight = basis_fnc.weight / wt


      # Rescale coefficients
      wt = zero
      for basis_fnc in self.dets:
         wt = wt + basis_fnc.weight
      for basis_fnc in self.dets:
         basis_fnc.weight = 4 * basis_fnc.weight / wt

      return copy.deepcopy(self.dets)

#-----------------------------------------------------------------------------

class Hubbard_Ham_1D:

   def __init__(self,n,t,U):
      # n = basis states
      self.n = n
      self.t = t
      self.U = U
      self.build_K_matrix()
      print "t =", self.t
      print "U =", self.U
      print "nsites =", self.n


   def build_K_matrix(self):
      # Hopping matrix
      n = self.n
      self.K = numpy.eye(n,n,1,dtype=numpy.float64)
      self.K = self.K + numpy.eye(n,n,-1,dtype=numpy.float64)
      self.K = self.K * (-self.t)
      #print "K =\n", self.K


   # Returns  <det0|H|det1> / <det0|det1>
   def eval_energy_via_G(self,det0,det1):
      (Ga, Gb) = det1.get_greens_func(det0)
      check_matrix_type([ Ga, Gb ])
      #print "Ga =\n", Ga
      #print "Gb =\n", Gb

      # One-electron terms
      # alpha
      Ka = zero
      for i in range(self.n):
         for j in range(self.n):
           Ka = Ka + self.K[i,j] * Ga[i,j]
      # beta
      Kb = zero
      for i in range(self.n):
         for j in range(self.n):
           Kb = Kb + self.K[i,j] * Gb[i,j]
      #print "\n< Ka > =", Ka
      #print "< Kb > =", Kb

      # Two-electron terms
      V = zero
      for i in range(self.n):
         V = V + self.U * Ga[i,i] * Gb[i,i]
      #print "< V >  =", V
      #print "< H >  =", V+Ka+Kb

      return V+Ka+Kb



#   # Returns  <det0|H|det1>
#   def eval_hamiltonian(self,det0,det1):
#
#      check_matrix_type([det0.phiA, det0.phiB, det1.phiA, det1.phiB])
#
#      # One-electron terms
#      (S, Sa, Sb) = det0.get_overlap_integrals(det1)
#      Ka = zero
#      Kb = zero
#      if det0.nalpha > 0:
#         Ka = numpy.linalg.det( det0.phiA.transpose() * self.K * det1.phiA)
#      if det0.nbeta > 0:
#         Kb = numpy.linalg.det( det0.phiB.transpose() * self.K * det1.phiB)
#      K = (Ka*Sb + Sa*Kb)
#
#      return K


#-----------------------------------------------------------------------------

class Trial_wavefunction:

   def __init__(self,n,nalpha,nbeta):
      self.det = Determinant(n,nalpha,nbeta)
#      self.det.phiA = numpy.zeros((n,nalpha),dtype=numpy.float64)
#      self.det.phiB = numpy.zeros((n,nbeta),dtype=numpy.float64)
      self.det.phiA = numpy.ones((n,nalpha),dtype=numpy.float64)
      self.det.phiB = numpy.ones((n,nbeta),dtype=numpy.float64)
      for i in range(nalpha):
         self.det.phiA[i,i] = one
      for i in range(nbeta):
         self.det.phiB[i+nalpha,i] = one
#         self.phiB[i,i] = one
      self.det.phiA = numpy.matrix(self.det.phiA)
      self.det.phiB = numpy.matrix(self.det.phiB)


#-----------------------------------------------------------------------------

class Determinant:

   def __init__(self,n,nalpha,nbeta):
      # n = basis states
      assert n >= nalpha
      assert n >= nbeta
      self.n = n
      self.nalpha = nalpha
      self.nbeta = nbeta
      self.init_phis()
      self.importance = None
      self.weight = one


   def __str__(self):
      print "phiA =\n", self.phiA
      print "phiB =\n", self.phiB
      print "Weight     =", self.weight
      print "Importance =", self.importance
      return ''


   def init_phis(self):
      self.phiA = numpy.zeros((self.n,self.nalpha),dtype=numpy.float64)
      self.phiB = numpy.zeros((self.n,self.nbeta),dtype=numpy.float64)
      self.phiA = numpy.matrix(self.phiA)
      self.phiB = numpy.matrix(self.phiB)


   def update_importance_fnc(self,det):
       self.importance = self.get_overlap_integrals(det)[0]


   def get_importance_fnc(self):
       assert (self.importance != None)
       return self.importance


   def apply_expM(self,Ma,Mb):
      # exp(M) |self>
      check_matrix_type([self.phiA, self.phiB])
      # alpha spin
      self.phiA = scipy.linalg.expm(Ma) * self.phiA
      # beta spin
      self.phiB = scipy.linalg.expm(Mb) * self.phiB
      # normalize
#      self.normalize()


   def get_greens_func(self,det0):
      # alpha spin
      Ga = numpy.zeros((self.n,self.n),dtype=numpy.float64)
      Ga = numpy.matrix(Ga)
      if self.nalpha > 0:
         check_matrix_type([ det0.phiA, self.phiA ])
         Ga = (det0.phiA.transpose() * self.phiA).getI()
         Ga = self.phiA * Ga * det0.phiA.transpose()
# Eqn.(7) ??
#         Ga = numpy.identity(self.n,dtype=numpy.float64) - Ga

      # beta spin
      Gb = numpy.zeros((self.n,self.n),dtype=numpy.float64)
      Gb = numpy.matrix(Gb)
      if self.nbeta > 0:
         check_matrix_type([ det0.phiB, self.phiB ])
         Gb = (det0.phiB.transpose() * self.phiB).getI()
         Gb = self.phiB * Gb * det0.phiB.transpose()
# Eqn.(7) ??
#         Gb = numpy.identity(self.n,dtype=numpy.float64) - Gb

      return Ga, Gb


   def get_overlap_integrals(self,det):
      # < det | self >
      check_matrix_type([self.phiA, self.phiB, det.phiA, det.phiB])
      alpha = one; beta = one
      if self.nalpha > 0:
         alpha = numpy.linalg.det( det.phiA.transpose() * self.phiA )
      if self.nbeta > 0:
         beta  = numpy.linalg.det( det.phiB.transpose() * self.phiB )
      return (alpha*beta, alpha, beta)


   def normalize(self):
      (ab,a,b) = self.get_overlap_integrals(self)
      norm = a**(one/(two*self.nalpha))
      self.phiA = self.phiA / norm
      norm = b**(one/(two*self.nbeta))
      self.phiB = self.phiB / norm


   def orthogonalize_orbitals(self):
      # Alpha
      self.phiA, R = numpy.linalg.qr(self.phiA)
      self.weight = self.weight * numpy.linalg.det( R )
      # Beta
      self.phiB, R = numpy.linalg.qr(self.phiB)
      self.weight = self.weight * numpy.linalg.det( R )


   def scale(self,fac):
      # Only scale alpha component for now
      self.phiA = self.phiA * fac


#-----------------------------------------------------------------------------

class HS_transformation:

   def __init__(self,n,U,dt):
      # n = basis states
      self.n = n
      self.U = U
      self.dt= dt
      self.gamma = numpy.arccosh(numpy.exp(half*self.dt*self.U))


   def get_one_site_HS(self,i,x):
      V = numpy.zeros((self.n,self.n),dtype=numpy.float64)
      V[i,i] = -half*self.dt*self.U + self.gamma*x
      return V


   def get_N_site_HS(self,site_and_phase):
      V = numpy.zeros((self.n,self.n),dtype=numpy.float64)
      for i in site_and_phase:
         V = V + self.get_one_site_HS(i[0],i[1])
      return V

#-----------------------------------------------------------------------------

class Projector:

   def __init__(self,ham,dt,Et,det0):
      self.ham = ham
      self.dt = dt
      self.Et = Et
      self.det0 = det0
      numpy.random.seed(1234567)

   def apply_expK(self,det):
      M = -self.dt*self.ham.K
      # alpha and beta have same propagator
      det.apply_expM(M,M)

   def apply_expKhalf(self,det):
      M = -(self.dt*half)*self.ham.K
      # alpha and beta have same propagator
      det.apply_expM(M,M)

   def apply_stochastic_expV(self,det):
      hst = HS_transformation(self.ham.n,self.ham.U,self.dt)
      # alpha spin 
      site_and_phase = []
      for i in range(self.ham.n):
         x = 2*numpy.random.randint(0,2) -1
         site_and_phase.append([i,x])
      Va = hst.get_N_site_HS(site_and_phase)
      # beta spin 
      for i in range(self.ham.n):
         site_and_phase[i][1] = -site_and_phase[i][1]
      Vb = hst.get_N_site_HS(site_and_phase)
      # propagate alpha and beta with same auxiliary fields
      det.apply_expM(Va,Vb)


   def apply_importance_sampled_stochastic_expV(self,det,det0):
      hst = HS_transformation(self.ham.n,self.ham.U,self.dt)
      # Loop over sites
      for i in range(self.ham.n):
         # Generate pair of dets
         Vp = hst.get_one_site_HS(i,1)
         Vm = hst.get_one_site_HS(i,-1)
         detP = copy.deepcopy(det)
         detM = copy.deepcopy(det)
         detP.apply_expM(Vp,Vm)
         detM.apply_expM(Vm,Vp)
         # Gather info about importance functions
         detP.update_importance_fnc(det0)
         detM.update_importance_fnc(det0)
         o0 = det.get_importance_fnc()
         op = detP.get_importance_fnc()
         om = detM.get_importance_fnc()
         # Select one det with adjusted probability
         Pp = op/o0 * half
         Pm = om/o0 * half
         N = Pp + Pm
         prob = Pp/N
         if (numpy.random.random() <= prob):
            det.apply_expM(Vp,Vm)
         else:
            det.apply_expM(Vm,Vp)
         # Update weight and importance function
         det.weight = det.weight * N
         det.update_importance_fnc(det0)


   def apply_deterministic_expV(self,det):
      assert (self.ham.n == det.n)
      hst = HS_transformation(self.ham.n,self.ham.U,self.dt)
      n64 = numpy.uint64(2**self.ham.n)
      dets = []

      # Apply 2**nsites HS propagators
      for k in range(n64):
         # Build bit string of auxiliary field variables
         bstr = numpy.zeros(self.ham.n,dtype=numpy.int8)
         j = 1
         foo = bin(k)[2:]
         foo = foo[::-1]
         for i in foo:
            bstr[self.ham.n-j] = i
            j = j+1
         # Apply this bit string of auxiliary fields to get 1 new det
         det_new = copy.deepcopy(det)

         # alpha spin 
         site_and_phase = []
         for i in range(self.ham.n):
            site_and_phase.append([i, 2*bstr[i]-1])
#
#
#         Vb = hst.get_N_site_HS(site_and_phase)
#         Va = hst.get_one_site_HS(0,1)
#         det_new.apply_expM(Va,Vb)
#         print Va
#         Va = hst.get_one_site_HS(1,1)
#         det_new.apply_expM(Va,Vb)
#         print Va
#         print det_new
#         sys.exit()
#
#
         Va = hst.get_N_site_HS(site_and_phase)
         # beta spin 
         for i in range(self.ham.n):
            site_and_phase[i][1] = -site_and_phase[i][1]
         Vb = hst.get_N_site_HS(site_and_phase)
         # propagate alpha and beta with same auxiliary fields
         det_new.apply_expM(Va,Vb)
         #print det_new
         #det_new.phiA = det_new.phiA * (half**self.ham.n)
         #det_new.phiB = det_new.phiB * (half**self.ham.n)
         dets.append(det_new)
      return dets


   def apply_stochastic_HS(self,det):
      # One-electron terms (first half transform)
      self.apply_expKhalf(det)
      det.update_importance_fnc(self.det0)

      # Two-electron terms
      self.apply_stochastic_expV(det)
#      self.apply_importance_sampled_stochastic_expV(det,det0)

      # One-electron terms (second half transform)
      self.apply_expKhalf(det)
      det.update_importance_fnc(self.det0)

      # Apply energy shift
#      det.weight = det.weight * numpy.exp(self.dt*self.Et)
      det.update_importance_fnc(self.det0)
      return [det]


   def apply_deterministic_HS(self,det):
##      # One-electron terms
##      self.apply_expK(det)
##      # Two-electron terms
##      dets = self.apply_deterministic_expV(det)
##      return dets

      # One-electron terms (first half transform)
      self.apply_expKhalf(det)
      # Two-electron terms
      dets = self.apply_deterministic_expV(det)
#      dets = [det]
      # One-electron terms (second half transform)
      for det in dets:
         self.apply_expKhalf(det)
      return dets
     

#-----------------------------------------------------------------------------

class Population:

   def orthogonalize_pop(self):
#      new_dets = []
#      for det0 in self.dets:
#         det0.normalize()
#         for det1 in new_dets:
#            s = det0.get_overlap_integrals(det1)[0]
#            det0.add(det1.scale(-s))
#         new_dets.append(det0)

      for det0 in self.dets:
         for det1 in self.dets:
            print "overlap =", det0.get_overlap_integrals(det1)[0]


   def get_pop_energy(self,ham):
##      num = zero
##      den = zero
##      for det in self.dets:
##         print "det energy =", ham.eval_energy_via_G(self.det0,det)
##         num = num + det.weight * ham.eval_energy_via_G(self.det0,det)
##         den = den + det.weight
##      Etot = num / den
##      return Etot, den

      overlaps = []
      for det in self.dets:
         overlaps.append(det.get_overlap_integrals(self.det0)[0])
      E = []
      for (det,overlap) in zip(self.dets,overlaps):
         e = det.weight * ham.eval_energy_via_G(self.det0,det)
         E.append(e * overlap)
      S = self.get_pop_overlap_with_psi(self.det0)
      return numpy.sum(E) / S, S

##
##      num = zero
##      for det in self.dets:
###         print ham.eval_hamiltonian(self.det0,det)
##         num = num + det.weight * ham.eval_hamiltonian(self.det0,det)
##      den = self.get_pop_overlap_with_psi(self.det0)
##      Etot = num / den
##      return Etot, den


   def get_pop_overlap_with_psi(self,det0):
      S = zero
      for det in self.dets:
         S = S + det.weight * det.get_overlap_integrals(det0)[0] 
      return S


   def propagate_time_step(self,projector,mode):
      new_dets = []
      for det in self.dets:
         if (mode == "stochastic"):
            new_dets = new_dets + projector.apply_stochastic_HS(det)
         elif (mode == "deterministic"):
            new_dets = new_dets + projector.apply_deterministic_HS(det)
         else:
            assert(False)
      return new_dets


   def rescale_pop_weights(self):
      wt = zero
      for det in self.dets:
         wt = wt + det.weight
      for det in self.dets:
         det.weight = det.weight / wt


#-----------------------------------------------------------------------------

class Stochastic_population(Population):

   def __init__(self,nsites,nalpha,nbeta,nwalkers):
      self.nalpha = nalpha
      self.nbeta = nbeta
      print "nalpha =", self.nalpha
      print "nbeta  =", self.nbeta
      # Trivial trial wavefunction
      trial_psi = Trial_wavefunction(nsites,nalpha,nbeta)
      self.det0 = trial_psi.det
      # Initialize stochastic population
      self.dets = []
      for n in range(nwalkers):
         det = copy.deepcopy(trial_psi.det)
         det.update_importance_fnc(self.det0)
         self.dets.append(det)
      self.nwalkers = len(self.dets)


   def evolve_pop(self,ham,dt,Et,nsteps,burnt):
      print "\nApplying propagator via sampled HS transform...\n"
      print "Number of walkers =", self.nwalkers
      print "\nInitial wavefunction =\n", self.det0
      projector = Projector(ham,dt,Et,self.det0)

      print '{0:18} {1:12} {2:12} {3:10}'.format('  tau','E','<E>','norm')
      E = []
      for n in range(1,nsteps):
         # Orthogonalize orbitals in each det
         if (n%10 == 0):
            for det in self.dets:
               det.orthogonalize_orbitals()
               det.update_importance_fnc(self.det0)
         # Propagate in imaginary time
         self.dets = self.propagate_time_step(projector,"stochastic")
         # Evaluate total energy
         (Etot, norm) = self.get_pop_energy(ham)
         if (n*dt > burnt): 
            E.append(Etot)
            if (n%1 == 0):
               print '{0:10.3f} {1:13.7f} {2:13.7f} {3:18.8f}'.format( n*dt, Etot, numpy.mean(E), norm)
         else:
            if (n%1 == 0): print '{0:10.3f} {1:12.8f}'.format(n*dt, Etot)
         # Rescale weights
         if (norm > self.nwalkers): self.rescale_pop_weights()



#-----------------------------------------------------------------------------

class Deterministic_population(Population):

   def __init__(self,nsites,nalpha,nbeta):
      self.nalpha = nalpha
      self.nbeta = nbeta
      print "nalpha =", self.nalpha
      print "nbeta  =", self.nbeta
      trial_psi = Trial_wavefunction(nsites,nalpha,nbeta)
      self.det0 = trial_psi.det
      self.dets = [self.det0]


   def evolve_pop(self,ham,dt,Et,nsteps):
      print "\nApplying propagator via exact HS transform...\n"
      print "\nTrial wavefunction =\n", self.det0
      projector = Projector(ham,dt,Et,self.det0)
      basis = Two_site_fci_basis()

      # Evaluate starting energy
      (Etot, norm) = self.get_pop_energy(ham)
      print '{0:10.3f}   E = {1:.8f}   (ndets={2:7d})     norm = {3:.8f}'.format(zero,Etot,1,norm)

      # Propagate in imaginary time
      for n in range(1,nsteps):
         # Apply HS-transformed projector
         self.dets = self.propagate_time_step(projector,"deterministic")
         # Project into FCI basis
         self.dets = basis.project_into_basis(self.dets)
         # Evaluate total energy
         (Etot, norm) = self.get_pop_energy(ham)
         print '{0:10.3f}   E = {1:.8f}   (ndets={2:7d})     norm = {3:.8f}'.format(n*dt,Etot,len(self.dets),norm)

      print

      
#-----------------------------------------------------------------------------

# 1-D HUBBARD MODEL

t = numpy.float64(1.0)
U = numpy.float64(1.0)
nsites = 2
nalpha = 1
nbeta  = 1
dt = numpy.float64(0.01)
print "dt =",dt
#Et = numpy.float64(-1.5616)
Et = numpy.float64(-0.0)

ham = Hubbard_Ham_1D(nsites,t,U)

pop = Deterministic_population(nsites,nalpha,nbeta)
pop.evolve_pop(ham,dt,Et,nsteps=10)

pop = Stochastic_population(nsites,nalpha,nbeta,nwalkers=10)
pop.evolve_pop(ham,dt,Et,nsteps=3,burnt=0.0)



